package top.hjcwzx.wrench.limiter.memory;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

//缓存策略
public abstract class ExpirableKeyStore<K, V> {

    Logger logger = LoggerFactory.getLogger(ExpirableKeyStore.class);

    public enum EvictionPolicy {
        LRU,
        LFU
    }

    private final EvictionPolicy evictionPolicy;

    // LRU 缓存，基于 Guava Cache
    private Cache<K, V> lruCache;

    // LFU 缓存和访问频率统计
    private Map<K, V> lfuCache;
    private Map<K, AtomicInteger> frequencyMap;

    // LFU 异步清理线程池
    private final ScheduledExecutorService cleaner = Executors.newSingleThreadScheduledExecutor();

    // 最大缓存大小
    private final int maxSize;

    /**
     * @param evictionPolicy      淘汰策略
     * @param expireAfterAccessMs 访问后过期时间（毫秒）
     * @param maxSize             最大缓存条目数
     */
    public ExpirableKeyStore(EvictionPolicy evictionPolicy, long expireAfterAccessMs, int maxSize) {
        this.evictionPolicy = evictionPolicy;
        this.maxSize = maxSize;

        if (evictionPolicy == EvictionPolicy.LRU) {
            this.lruCache = CacheBuilder.newBuilder()
                    .maximumSize(maxSize)
                    .expireAfterAccess(expireAfterAccessMs, TimeUnit.MILLISECONDS)
                    .build();
        } else if (evictionPolicy == EvictionPolicy.LFU) {
            this.lfuCache = new ConcurrentHashMap<>();
            this.frequencyMap = new ConcurrentHashMap<>();

            // 每分钟清理一次频率最低的部分缓存
            cleaner.scheduleAtFixedRate(this::cleanLFUCache, 1, 1, TimeUnit.MINUTES);
        }
    }

    /**
     * 获取缓存值，访问时更新访问频率（LFU）
     */
    public V get(K key) {
        if (evictionPolicy == EvictionPolicy.LRU) {
            return lruCache.getIfPresent(key);
        } else {
            AtomicInteger freq = frequencyMap.get(key);
            if (freq != null) {
                freq.incrementAndGet();
            }
            return lfuCache.get(key);
        }
    }

    /**
     * 放入缓存，LFU 新增时默认频率为1
     */
    public void put(K key, V value) {
        if (evictionPolicy == EvictionPolicy.LRU) {
            lruCache.put(key, value);
        } else {
            lfuCache.put(key, value);
            frequencyMap.putIfAbsent(key, new AtomicInteger(1));
            // 如果超过最大容量，异步清理也会处理
            if (lfuCache.size() > maxSize) {
                cleaner.submit(this::cleanLFUCache);
            }
        }
    }

    /**
     * 移除缓存
     */
    public void remove(K key) {
        if (evictionPolicy == EvictionPolicy.LRU) {
            lruCache.invalidate(key);
        } else {
            lfuCache.remove(key);
            frequencyMap.remove(key);
        }
    }

    /**
     * LFU 异步清理频率最低的缓存，剔除10%的缓存条目（可调）
     */
    private void cleanLFUCache() {
        try {
            int cacheSize = lfuCache.size();
            if (cacheSize <= maxSize) {
                return;
            }
            // 清理10%
            int removeCount = cacheSize / 10;

            // 按频率排序，取最低的 keys
            List<K> keysToRemove = frequencyMap.entrySet().stream()
                    .sorted((e1, e2) -> Integer.compare(e1.getValue().get(), e2.getValue().get()))
                    .limit(removeCount)
                    .map(Map.Entry::getKey)
                    .collect(Collectors.toList());

            for (K key : keysToRemove) {
                lfuCache.remove(key);
                frequencyMap.remove(key);
                logger.debug("limiter - LFU内存淘汰策略触发 - key -" + key);
            }
        } catch (Exception e) {
            // 日志异常，这里简单打印
            e.printStackTrace();
        }
    }

    /**
     * 关闭清理线程
     */
    public void shutdown() {
        cleaner.shutdown();
    }
}
